module poly_cubic_mod
  use recon_type_mod
  use recon_math_mod
  use poly_utils_mod

  implicit none
  
contains

  subroutine poly_cubic_init(this, nd, sw, swx, swy, swz, xc, yc, zc, dx, dy, dz, ic, jc, kc, is, ie, js, je, ks, ke, mask, id)

    class(recon_type), intent(inout) :: this
    integer , intent(in), optional :: nd
    integer , intent(in), optional :: sw
    integer , intent(in), optional :: swx
    integer , intent(in), optional :: swy
    integer , intent(in), optional :: swz
    real(16), intent(in), optional :: xc(:)
    real(16), intent(in), optional :: yc(:)
    real(16), intent(in), optional :: zc(:)
    real(16), intent(in), optional :: dx
    real(16), intent(in), optional :: dy
    real(16), intent(in), optional :: dz
    integer , intent(in), optional :: ic
    integer , intent(in), optional :: jc
    integer , intent(in), optional :: kc
    integer , intent(in), optional :: is
    integer , intent(in), optional :: ie
    integer , intent(in), optional :: js
    integer , intent(in), optional :: je
    integer , intent(in), optional :: ks
    integer , intent(in), optional :: ke
    integer , intent(in), optional :: mask(:,:,:) ! Cell mask
    integer , intent(in), optional :: id
    
    integer i, j, k

    call this%clear()

    if(nd/=3)then
      stop 'poly_cubic_init support only nd==3'
    endif
    
    this%sw  = sw
    if(present(swx))then
      this%swx = swx
    else
      this%swx = this%sw
    endif
    if(present(swy))then
      this%swy = swy
    else
      this%swy = this%sw
    endif
    if(present(swz))then
      this%swz = swz
    else
      this%swz = this%sw
    endif
    this%nd  = nd
    this%nc  = this%swx * this%swy * this%swz
    this%rw  = ( this%sw - 1 ) / 2

    if(present(ic))then
      this%ic = ic
    else
      this%ic = this%rw + 1
    endif

    if(present(jc))then
      this%jc = jc
    else
      this%jc = this%rw + 1
    endif

    if(present(kc))then
      this%kc = kc
    else
      this%kc = this%rw + 1
    endif
    
    if(present(dx))this%dx = dx
    if(present(dy))this%dy = dy
    if(present(dz))this%dz = dz

    if(present(is))then; this%is = is; else; this%is = 1       ; endif
    if(present(ie))then; this%ie = ie; else; this%ie = this%swx; endif
    if(present(js))then; this%js = js; else; this%js = 1       ; endif
    if(present(je))then; this%je = je; else; this%je = this%swy; endif
    if(present(ks))then; this%ks = ks; else; this%ks = 1       ; endif
    if(present(ke))then; this%ke = ke; else; this%ke = this%swz; endif

    allocate(this%cell_mask(this%is:this%ie,this%js:this%je,this%ks:this%ke)); this%cell_mask = 0
    allocate(this%poly_mask(this%swx,this%swy,this%swz)); this%poly_mask = 0

    if (present(mask)) then
      this%cell_mask = mask
    else
      this%cell_mask = 1
    end if

    this%ac = count(this%cell_mask == 1)

    this%rwL = this%rw
    this%rwR = this%rw
    this%rwB = this%rw
    this%rwT = this%rw
    this%rwD = this%rw
    this%rwU = this%rw
    do i = 0, -this%rw, -1
      if( this%cell_mask(this%ic+i,1,1) == 1 )then
        this%rwL = min(abs(i), this%rw)
      endif
    enddo
    do i = 0, this%rw
      if( this%cell_mask(this%ic+i,1,1) == 1 )then
        this%rwR = min(abs(i), this%rw)
      endif
    enddo

    do j = 0, -this%rw, -1
      if( this%cell_mask(1,this%jc+j,1) == 1 )then
        this%rwB = min(abs(j), this%rw)
      endif
    enddo
    do j = 0, this%rw
      if( this%cell_mask(1,this%jc+j,1) == 1 )then
        this%rwT = min(abs(j), this%rw)
      endif
    enddo

    do k = 0, -this%rw, -1
      if( this%cell_mask(1,1,this%kc+k) == 1 )then
        this%rwD = min(abs(k), this%rw)
      endif
    enddo
    do k = 0, this%rw
      if( this%cell_mask(1,1,this%kc+k) == 1 )then
        this%rwU = min(abs(k), this%rw)
      endif
    enddo

    ! Set polynomial mask from cell mask.
    this%poly_mask(:,:,1:this%rwU+this%rwD+1) = 1

    allocate(this%xc(this%is:this%ie))
    allocate(this%yc(this%js:this%je))
    allocate(this%zc(this%ks:this%ke))

    ! Set coordinates of cells on the large stencil with origin at center.
    do i = -this%rwL, this%rwR
      this%xc(this%ic+i) = i
    end do
    do j = -this%rwB, this%rwT
      this%yc(this%jc+j) = j
    end do
    do k = -this%rwD, this%rwU
      this%zc(this%kc+k) = k
    end do

    this%initialized = .true.

  end subroutine poly_cubic_init

  subroutine poly_cubic_calc_recon_matrix(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    real(16), allocatable, dimension(:,:) :: A, iA
    integer i, j, k, ipt, n, iterm, dxn, dyn, dzn

    if (allocated(this%iA            )) deallocate(this%iA            )
    if (allocated(this%poly          )) deallocate(this%poly          )
    if (allocated(this%poly_r16      )) deallocate(this%poly_r16      )
    if (allocated(this%recon_mtx     )) deallocate(this%recon_mtx     )
    if (allocated(this%recon_mtx_r16 )) deallocate(this%recon_mtx_r16 )
    if (allocated(this%dpoly         )) deallocate(this%dpoly         )
    if (allocated(this%dpoly_r16     )) deallocate(this%dpoly_r16     )
    if (allocated(this%drecon_mtx    )) deallocate(this%drecon_mtx    )
    if (allocated(this%drecon_mtx_r16)) deallocate(this%drecon_mtx_r16)

    allocate(this%iA            (this%nc ,this%nc));                                        this%iA             = 0
    allocate(this%poly          (this%npt,this%nc));                                        this%poly           = 0
    allocate(this%poly_r16      (this%npt,this%nc));                                        this%poly_r16       = 0
    allocate(this%recon_mtx     (this%npt,this%nc));                                        this%recon_mtx      = 0
    allocate(this%recon_mtx_r16 (this%npt,this%nc));                                        this%recon_mtx_r16  = 0
    allocate(this%dpoly         (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%dpoly          = 0
    allocate(this%dpoly_r16     (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%dpoly_r16      = 0
    allocate(this%drecon_mtx    (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%drecon_mtx     = 0
    allocate(this%drecon_mtx_r16(this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%drecon_mtx_r16 = 0

    allocate( A(this%nc,this%nc)); A  = 0
    allocate(iA(this%nc,this%nc)); iA = 0

    ! Set the poly for each evaluation point.
    ! Select monomials according to mask.
    do ipt = 1, this%npt
      iterm = 0
      do k = 1, this%swz
        do j = 1, this%swy
          do i = 1, this%swx
            if (this%poly_mask(i,j,k) == 1) then
              iterm = iterm + 1
              call calc_monomial(this%x(ipt), this%y(ipt), this%z(ipt), i - 1, j - 1, k - 1, this%poly_r16(ipt,iterm))
              do dzn = 0, this%swz - 1
                do dyn = 0, this%swy - 1
                  do dxn = 0, this%swx - 1
                    call calc_deriv_monomial( this%x(ipt), this%y(ipt), this%z(ipt), &
                                              i - 1, j - 1, k - 1,                   &
                                              dxn, dyn, dzn,                         &
                                              this%swx, this%swy, this%swz,          &
                                              this%dpoly_r16(ipt,iterm,dxn,dyn,dzn) )
                    this%dpoly_r16(ipt,iterm,dxn,dyn,dzn) = this%dpoly_r16(ipt,iterm,dxn,dyn,dzn) / ( this%dx**dxn * this%dy**dyn * this%dz**dzn )
                  end do
                end do
              end do
            endif
          end do
        end do
      end do
    end do
    this%poly  = this%poly_r16
    this%dpoly = this%dpoly_r16
    
    ! Calculate inverse of integral coefficient matrix.
    n = this%ac
    call calc_tensor_product_poly_integral_matrix(this%swx, this%swy, this%swz, this%xc, this%yc, this%zc, A, this%cell_mask, this%poly_mask)

    call inverse_matrix(A(1:n,1:n), iA(1:n,1:n), ierr)
    if (ierr /= 0) then
      deallocate(A, iA)
      return
    end if
    
    this%recon_mtx_r16(:,1:n) = matmul(this%poly_r16(:,1:n), iA(1:n,1:n))
    
    do dzn = 0, this%swz - 1
      do dyn = 0, this%swy - 1
        do dxn = 0, this%swx - 1
          this%drecon_mtx_r16(:,1:n,dxn,dyn,dzn) = matmul(this%dpoly_r16(:,1:n,dxn,dyn,dzn), iA(1:n,1:n))
        end do
      end do
    end do

    this%iA         = iA                 
    this%recon_mtx  = this%recon_mtx_r16 
    this%drecon_mtx = this%drecon_mtx_r16
    
    deallocate(A, iA)

  end subroutine poly_cubic_calc_recon_matrix

  subroutine poly_cubic_calc_poly(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision
    integer i, j, k, ipt, n, iterm, dxn, dyn, dzn
    
    ierr = 0

    if (allocated(this%poly    )) deallocate(this%poly    )
    if (allocated(this%poly_r16)) deallocate(this%poly_r16)

    allocate(this%poly     (this%npt,this%nc)); this%poly     = 0
    allocate(this%poly_r16 (this%npt,this%nc)); this%poly_r16 = 0

    ! Set the poly for each evaluation point.
    ! Select monomials according to mask.
    do ipt = 1, this%npt
      iterm = 0
      do k = 1, this%swz
        do j = 1, this%swy
          do i = 1, this%swx
            if (this%poly_mask(i,j,k) == 1) then
              iterm = iterm + 1
              call calc_monomial(this%x(ipt), this%y(ipt), this%z(ipt), i - 1, j - 1, k - 1, this%poly_r16(ipt,iterm))
            endif
          end do
        end do
      end do
    end do
    this%poly = this%poly_r16

  end subroutine poly_cubic_calc_poly

  subroutine poly_cubic_calc_dpoly(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    integer i, j, k, ipt, n, iterm, dxn, dyn, dzn

    if (allocated(this%dpoly         )) deallocate(this%dpoly         )
    if (allocated(this%dpoly_r16     )) deallocate(this%dpoly_r16     )

    allocate(this%dpoly         (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%dpoly          = 0
    allocate(this%dpoly_r16     (this%npt,this%nc,0:this%swx-1,0:this%swy-1,0:this%swz-1)); this%dpoly_r16      = 0

    ierr = 0
    
    ! Set the poly for each evaluation point.
    ! Select monomials according to mask.
    do ipt = 1, this%npt
      iterm = 0
      do k = 1, this%swz
        do j = 1, this%swy
          do i = 1, this%swx
            if (this%poly_mask(i,j,k) == 1) then
              iterm = iterm + 1
              do dzn = 0, this%swz - 1
                do dyn = 0, this%swy - 1
                  do dxn = 0, this%swx - 1
                    call calc_deriv_monomial( this%x(ipt), this%y(ipt), this%z(ipt), &
                                              i - 1, j - 1, k - 1,                   &
                                              dxn, dyn, dzn,                         &
                                              this%swx, this%swy, this%swz,          &
                                              this%dpoly_r16(ipt,iterm,dxn,dyn,dzn) )
                    this%dpoly_r16(ipt,iterm,dxn,dyn,dzn) = this%dpoly_r16(ipt,iterm,dxn,dyn,dzn) / ( this%dx**dxn * this%dy**dyn * this%dz**dzn )
                  end do
                end do
              end do
            endif
          end do
        end do
      end do
    end do
    this%dpoly = this%dpoly_r16

  end subroutine poly_cubic_calc_dpoly

  subroutine poly_cubic_calc_iA(this, ierr)

    class(recon_type), intent(inout) :: this
    integer, intent(out) :: ierr

    ! Local double double arrays for preserving precision.
    real(16), allocatable, dimension(:,:) :: A, iA
    integer i, j, k, ipt, n, iterm, dxn, dyn, dzn

    if (allocated(this%iA)) deallocate(this%iA)

    allocate(this%iA(this%nc ,this%nc));this%iA = 0

    allocate( A(this%nc,this%nc)); A  = 0
    allocate(iA(this%nc,this%nc)); iA = 0
    
    ! Calculate inverse of integral coefficient matrix.
    n = this%ac
    call calc_tensor_product_poly_integral_matrix(this%swx, this%swy, this%swz, this%xc, this%yc, this%zc, A, this%cell_mask, this%poly_mask)

    call inverse_matrix(A(1:n,1:n), iA(1:n,1:n), ierr)
    if (ierr /= 0) then
      deallocate(A, iA)
      return
    end if

    this%iA = iA                 
    
    deallocate(A, iA)

  end subroutine poly_cubic_calc_iA
  
  subroutine poly_cubic_reconstruct(this, fi, fo, TCI, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points
    logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fo = matmul(this%recon_mtx(:,1:this%ac), pack(fi, .true.))
    
  end subroutine poly_cubic_reconstruct
  
  subroutine poly_cubic_reconstruct_by_poly(this, poly, fi, fo, TCI, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: poly(:,:) ! Polynomial with dimension of (npt,this%ac)
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values, total length must be equal to this%ac
    real(8), intent(out) :: fo(:)     ! Reconstructed function values on evaluation points, dimension (npt)
    logical, intent(in ), optional :: TCI ! Trouble Cell Indicator(TCI), 1 for existing TC(use WENO), 0 for no TC(use poly)
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fo = matmul( poly(:,1:this%ac), matmul( this%iA(1:this%ac,1:this%ac) , pack(fi, .true.) ) )
    
  end subroutine poly_cubic_reconstruct_by_poly
  
  subroutine poly_cubic_reconstruct_deriv(this, fi, fx, fy, fz, ierr)

    class(recon_type), intent(inout) :: this
    real(8), intent(in ) :: fi(:,:,:) ! Cell averaged function values
    real(8), intent(out), optional :: fx(:)     ! Reconstructed x derivative values on evaluation points
    real(8), intent(out), optional :: fy(:)     ! Reconstructed y derivative values on evaluation points
    real(8), intent(out), optional :: fz(:)     ! Reconstructed y derivative values on evaluation points
    integer, intent(out), optional :: ierr

#ifdef DEBUG
    if(present(ierr))then
      ierr = 0
      if (size(fi) /= this%nc) then
        if (present(ierr)) ierr = 1
        return
      end if
      if (size(fo) /= this%npt) then
        if (present(ierr)) ierr = 1
        return
      end if
    endif
#endif

    fx = matmul(this%drecon_mtx(:,1:this%ac,1,0,0), pack(fi, .true.))
    fy = matmul(this%drecon_mtx(:,1:this%ac,0,1,0), pack(fi, .true.))
    fz = matmul(this%drecon_mtx(:,1:this%ac,0,0,1), pack(fi, .true.))
    
  end subroutine poly_cubic_reconstruct_deriv

end module poly_cubic_mod
